import pygame
import time
import json
import os
import numpy as np
from scipy import stats
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences

pred_width = 250
USER_STATS_PATH = 'model/precise_model/user_stats.json'
class TextBox:
    def __init__(self, screen, font, example_text=None, text_file=None, username=None, is_existing_user=False):
        self.screen = screen
        self.font = font
        self.example_text = example_text
        if text_file:
            self.example_text = self.load_game_texts(text_file)
        self.input_text = ""
        self.correct_color = (0, 255, 0)  # 正确字符颜色（绿色）
        self.error_color = (255, 0, 0)    # 错误字符颜色（红色）
        self.default_color = (0, 0, 0)    # 默认字符颜色（黑色）
        self.scroll_offset = 0
        self.locked_indices = set()
        self.error_index = -1
        self.username = username
        
        # 初始化按键记录相关属性
        self.key_events = []
        self.start_time = None
        self.first_key_pressed = False
        self.group_start_time = None
        self.valid_key_count = 0
        self.sequence_length = 0
        
        # 计时器暂停相关属性
        self.timer_paused = False
        self.pause_start_time = 0
        self.paused_duration = 0
        
        # 用户类型判断和区域划分
        self.is_existing_user = is_existing_user
        self.update_text_area()  # 初始化文本区域
        self.midpoint_position = 0  # 记录中点位置
        
        # 初始化模型相关
        self.model = None
        self.key_to_index = {}
        self.label_to_index = {}
        self.index_to_label = {}
        if self.is_existing_user:  # 只有已有用户才初始化模型
            self.initialize_model()
        
        # 预测相关
        self.last_prediction = None
        self.current_prediction = None
        self.char_width_cache = {}

        self.midpoint = len(self.example_text) // 2  # 整数除法
        self.has_passed_midpoint = False  # 是否已过中点
        self.all_user_probs_history = []  # 存储每次预测的完整概率分布

    def update_text_area(self):
        """更新文本区域"""
        screen_width = self.screen.get_width()
        
        # 新用户或小屏幕：文本框占满底部
        if not self.is_existing_user or screen_width < 1000:
            self.text_area_width = screen_width
            self.box_rect = pygame.Rect(
                0, 
                int(self.screen.get_height() * 0.90),
                screen_width,
                int(self.screen.get_height() * 0.10)
            )
        # 已有用户大屏幕：右侧留出预测区域
        else:
            self.text_area_width = screen_width - pred_width
            self.box_rect = pygame.Rect(
                0,
                int(self.screen.get_height() * 0.90),
                self.text_area_width,
                int(self.screen.get_height() * 0.10)
            )

    @staticmethod
    def load_game_texts(filepath):
        """加载游戏语料，删除换行符并将所有文本合并为一行"""
        with open(filepath, 'r', encoding='utf-8') as file:
            lines = [line.strip() for line in file if line.strip()]
            return " ".join(lines)

    def initialize_model(self):
        """初始化生物特征识别模型"""
        try:
            # 使用正确的模型路径（注意第一个脚本保存的是.keras格式）
            model_path = 'model/precise_model/lstm_model.keras'  # 修改为.keras扩展名
            label_path = 'model/precise_model/lstm_label_to_index.json'
            
            if not os.path.exists(model_path):
                raise FileNotFoundError(f"模型文件不存在: {model_path}")
                
            self.model = load_model(model_path)
            
            # 加载标签映射
            with open(label_path, 'r') as f:
                self.label_to_index = json.load(f)
            self.index_to_label = {v: k for k, v in self.label_to_index.items()}
            
            print("生物特征模型加载成功")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            self.model = None

    def start_timer(self):
        """开始计时，记录游戏开始时间"""
        self.start_time = time.time()

    def handle_event(self, event):
        """处理事件（兼容handle_input和handle_resize）"""
        if event.type == pygame.VIDEORESIZE:
            self.handle_resize(event)
        else:
            self.handle_input(event)

    def handle_resize(self, event):
        """处理窗口大小变化"""
        self.update_text_area()
        self.char_width_cache.clear()
        
        # 窗口大小改变时重新计算可见字符和滚动位置
        self.calculate_visible_chars()
        self.update_scroll_offset()

    def get_char_width(self, char):
        """获取字符宽度，使用缓存提高性能"""
        if char not in self.char_width_cache:
            self.char_width_cache[char] = self.font.size(char)[0]
        return self.char_width_cache[char]

    def render(self):
        """渲染文本框和内容"""
        # 绘制文本框
        pygame.draw.rect(self.screen, (0, 0, 0), self.box_rect, 2)

        # 每次渲染前更新显示区域和滚动位置
        self.calculate_visible_chars()
        self.update_scroll_offset()

        # 已有用户需要绘制预测区域
        if self.is_existing_user:
            self.render_prediction_area()
        
        # 渲染文本内容
        self.update_scroll_offset()
        self.render_example_text()
        self.render_input_text()

    def render_prediction_area(self):
        """渲染预测区域（仅对已有用户）"""
        # 计算屏幕宽度
        screen_width = self.screen.get_width()
        
        # 小屏幕不显示预测区域（小于800像素）
        if screen_width < 1000:
            return
        
        pred_rect = pygame.Rect(
            self.box_rect.right,  # 从文本框右侧开始
            self.box_rect.y,
            pred_width,  # 固定宽度
            self.box_rect.height
        )
        
        # 确保预测区域不会超出屏幕
        if pred_rect.right > screen_width:
            pred_rect.width = screen_width - pred_rect.x
        
        pygame.draw.rect(self.screen, (240, 240, 240), pred_rect)  # 浅灰色背景
        pygame.draw.rect(self.screen, (0, 0, 0), pred_rect, 2)  # 边框
        
        # 预测内容渲染
        if self.current_prediction or self.last_prediction:
            try:
                pred_font = pygame.font.Font('assets/fonts/KeLaiYinLanDeYeWan-2.ttf', 20)
            except:
                pred_font = pygame.font.Font(None, 20)
            
            # 标题
            title = pred_font.render("预测结果:", True, (0, 0, 0))
            self.screen.blit(title, (pred_rect.x + 5, pred_rect.y + 5))
            
            y_offset = 30
            # 当前预测
            if self.current_prediction:
                curr_text = f"当前预测: {self.current_prediction[0]} {self.current_prediction[1]*100:.0f}%"
                color = (0, 180, 0) if self.current_prediction[1] >= 0.8 else (180, 0, 0)
                curr_surface = pred_font.render(curr_text, True, color)
                self.screen.blit(curr_surface, (pred_rect.x + 5, pred_rect.y + y_offset))
                y_offset += 25
            
            # 上次预测
            if self.last_prediction:
                last_text = f"上次预测: {self.last_prediction[0]} {self.last_prediction[1]*100:.0f}%"
                color = (0, 150, 0) if self.last_prediction[1] >= 0.8 else (150, 0, 0)
                last_surface = pred_font.render(last_text, True, color)
                self.screen.blit(last_surface, (pred_rect.x + 5, pred_rect.y + y_offset))

    def update_scroll_offset(self):
        """更新滚动偏移量，基于当前输入位置与显示区域的关系"""
        if not self.input_text or not hasattr(self, 'box_rect'):
            self.scroll_offset = 0
            self.has_passed_midpoint = False
            return

        # 计算显示区域能容纳的字符数（动态适应窗口大小）
        self.calculate_visible_chars()
        
        # 当前输入位置（最后一个字符的索引）
        current_pos = len(self.input_text) - 1
        
        # 显示区域中点位置（取可见字符数的一半）
        midpoint = self.visible_chars // 2
        
        # 如果当前输入位置超过中点且还未开始滚动，则标记需要滚动
        if current_pos > midpoint and not self.has_passed_midpoint:
            self.has_passed_midpoint = True
        
        # 如果需要滚动，计算偏移量使当前输入位置保持在显示区域中点
        if self.has_passed_midpoint:
            self.scroll_offset = max(0, current_pos - midpoint)


    def calculate_visible_chars(self):
        """动态计算当前显示区域能容纳的字符数"""
        if not hasattr(self, 'box_rect') or not hasattr(self, 'example_text'):
            self.visible_chars = 0
            return
        
        x = self.box_rect.x + 10
        right_boundary = self.box_rect.x + self.text_area_width - 10
        count = 0
        
        # 从当前滚动位置开始计算
        text_to_check = self.example_text[self.scroll_offset:]
        
        for char in text_to_check:
            char_width = self.get_char_width(char)
            if x + char_width > right_boundary:
                break
            x += char_width
            count += 1
        
        self.visible_chars = count

    def render_example_text(self):
        """渲染示例文本"""
        x = self.box_rect.x + 10
        y = self.box_rect.y + 10
        right_boundary = self.box_rect.x + self.text_area_width - 10

        # 应用滚动偏移
        visible_text = self.example_text[self.scroll_offset:]
        
        for i, char in enumerate(visible_text):
            char_width = self.get_char_width(char)
            if x + char_width > right_boundary:
                break

        visible_text = self.example_text[self.scroll_offset:]

        for i, char in enumerate(visible_text):
            char_width = self.get_char_width(char)
            if x + char_width > right_boundary:
                break

            color = self.default_color
            global_index = self.scroll_offset + i
            
            if global_index < len(self.input_text) and global_index not in self.locked_indices:
                if self.input_text[global_index] == char:
                    color = self.correct_color
                    self.locked_indices.add(global_index)
                elif global_index == self.error_index:
                    color = self.error_color

            char_surface = self.font.render(char, True, color)
            self.screen.blit(char_surface, (x, y))
            x += char_width

    def render_input_text(self):
        """动态渲染输入文本，精确计算每个字符位置"""
        x = self.box_rect.x + 10
        y = self.box_rect.y + 40
        right_boundary = self.box_rect.x + self.text_area_width - 10

        visible_input = self.input_text[self.scroll_offset:]

        for i, char in enumerate(visible_input):
            char_width = self.get_char_width(char)
            if x + char_width > right_boundary:
                break

            global_index = self.scroll_offset + i
            color = self.default_color
            
            if global_index in self.locked_indices:
                color = self.correct_color
            elif global_index == self.error_index:
                color = self.error_color

            char_surface = self.font.render(char, True, color)
            self.screen.blit(char_surface, (x, y))
            x += char_width

    def handle_input(self, event):
        """处理用户输入，并记录按键信息"""
        if event.type == pygame.KEYDOWN:
            # 如果是第一次按键，记录绝对开始时间
            if not self.first_key_pressed:
                self.first_key_pressed = True
                self.start_time = time.time()
                self.absolute_start_time = time.time()  # 新增：记录绝对开始时间

            # 处理退格键
            if event.key == pygame.K_BACKSPACE:
                if len(self.input_text) > 0 and (len(self.input_text) - 1) not in self.locked_indices:
                    self.input_text = self.input_text[:-1]
                    if self.error_index == len(self.input_text):
                        self.error_index = -1
            
            # 处理普通字符
            elif event.unicode and (event.unicode.isalnum() or event.unicode.isspace() or event.unicode in ',.'):
                if self.error_index == -1:  # 只有在没有错误时才允许输入
                    current_absolute_time = time.time()  # 获取当前绝对时间
                    
                    # 记录按键按下事件（使用绝对时间）
                    key_event = {
                        "key_content": event.unicode,
                        "absolute_press_time": current_absolute_time - self.absolute_start_time,
                        "absolute_release_time": None,  # 将在KEYUP时设置
                        "is_correct": None
                    }
                    self.key_events.append(key_event)

                    # 更新用户输入文本
                    self.input_text += event.unicode
                    if len(self.input_text) <= len(self.example_text):
                        if self.input_text[-1] == self.example_text[len(self.input_text) - 1]:
                            if hasattr(self, 'game') and hasattr(self.game, 'hero'):
                                self.game.hero.add_move_pixels(self.game.hero.char_move_speed)
                                self.game.hero.update()
                            key_event["is_correct"] = True
                        else:
                            key_event["is_correct"] = False
                            self.error_index = len(self.input_text) - 1

                    self.valid_key_count += 1
                    
                    # 修改预测调用部分
                    if (self.valid_key_count) % 10 == 1 and self.valid_key_count > 1 and self.model is not None:
                        # 确保有足够的数据（至少10个完整击键）
                        if len([e for e in self.key_events[-11:-1] if e["absolute_release_time"] is not None]) >= 10:
                            result = self.predict_identity(self.key_events[-11:-1])
                            if result:
                                predicted_user, confidence = result
                                print(f"预测结果: 用户 {predicted_user}, 置信度 {confidence:.2f}")

        elif event.type == pygame.KEYUP:
            current_absolute_time = time.time()
            # 找到所有未记录的相同按键事件
            unmatched_events = [
                e for e in self.key_events 
                if e["key_content"] == event.unicode and e["absolute_release_time"] is None
            ]
            
            if unmatched_events:
                # 使用最早的未匹配按键事件（FIFO原则）
                key_event = unmatched_events[0]
                key_event["absolute_release_time"] = current_absolute_time - self.absolute_start_time
                
    def predict_identity(self, keystrokes):
        """使用与LSTM训练完全一致的特征提取方式进行预测"""
        try:
            if not keystrokes or len(keystrokes) < 10:
                print("预测失败: 需要至少10次击键数据")
                return None

            # 1. 加载用户统计信息
            with open(USER_STATS_PATH) as f:
                loaded_stats = json.load(f)  # 改名为 loaded_stats 避免命名冲突
                user_stats = loaded_stats['user_stats']
                global_stats = loaded_stats['global_stats']

            # 2. 计算当前序列的统计量（使用中位数与训练时一致）
            durations = []
            intervals = []
            for i in range(len(keystrokes)):
                duration = keystrokes[i]['absolute_release_time'] - keystrokes[i]['absolute_press_time']
                durations.append(duration)
                if i > 0:
                    interval = keystrokes[i]['absolute_press_time'] - keystrokes[i-1]['absolute_press_time']
                    intervals.append(interval)
            
            current_median_duration = np.median(durations) if durations else global_stats['durations'][0]
            current_median_interval = np.median(intervals) if intervals else global_stats['intervals'][0]

            # 3. 构建特征序列（与训练完全相同的11维特征）
            sequence = []
            for i in range(len(keystrokes)):
                k = keystrokes[i]
                duration = k['absolute_release_time'] - k['absolute_press_time']
                rel_duration = duration / (current_median_duration + 1e-6)
                
                # 初始化特征向量（全零）
                features = np.zeros(11, dtype=np.float32)  # 注意改为11维
                
                # 基础特征（始终存在）
                features[0] = duration
                features[1] = rel_duration
                features[10] = i / len(keystrokes)  # 位置信息
                
                if i > 0:
                    prev = keystrokes[i-1]
                    interval = k['absolute_press_time'] - prev['absolute_press_time']
                    rel_interval = interval / (current_median_interval + 1e-6)
                    prev_duration = prev['absolute_release_time'] - prev['absolute_press_time']
                    accel = (duration - prev_duration) / (interval + 1e-6)
                    
                    # Bigram处理
                    bigram = (prev['key_content'], k['key_content'])
                    bigram_str = f"{bigram[0]}_{bigram[1]}"
                    
                    # 查找用户统计信息中的bigram数据
                    bigram_features = [0.5, 0.5, 1.0, 1.0]  # 默认中性值
                    for user, data in user_stats.items():
                        if bigram_str in data['bigrams']:
                            samples = data['bigrams'][bigram_str]
                            if len(samples) >= 3:  # 与训练时相同的阈值
                                median_duration = np.median([x[0] for x in samples])
                                median_interval = np.median([x[1] for x in samples])
                                bigram_features = [
                                    stats.percentileofscore([x[0] for x in samples], duration) / 100.0,  # 使用 scipy.stats
                                    stats.percentileofscore([x[1] for x in samples], interval) / 100.0,   # 使用 scipy.stats
                                    duration / (median_duration + 1e-6),
                                    interval / (median_interval + 1e-6)
                                ]
                                break
                    
                    features[2] = interval
                    features[3] = rel_interval
                    features[4] = accel
                    features[5:9] = bigram_features
                
                sequence.append(features)

             # --- 模型预测 ---
            padded_sequence = pad_sequences([sequence], maxlen=100, padding='post', dtype='float32')
            all_user_probs = self.model.predict(padded_sequence, verbose=0)[0]  # 所有用户的概率分布

            # 存储当前预测的所有用户概率
            self.all_user_probs_history.append(all_user_probs)

            # 如果不足5次预测，不输出结果
            if len(self.all_user_probs_history) < 5:
                return None

            # --- 综合最近5次预测 ---
            # 初始化总概率字典 {用户索引: 总概率}
            total_probs = {i: 0.0 for i in range(len(all_user_probs))}

            # 累加5次预测的概率
            for probs in self.all_user_probs_history[-5:]:
                for user_idx, prob in enumerate(probs):
                    total_probs[user_idx] += prob

            # 计算每个用户的平均概率
            avg_probs = {
                user_idx: total_prob / 5 
                for user_idx, total_prob in total_probs.items()
            }

            # 按平均概率从高到低排序
            sorted_avg_probs = sorted(
                avg_probs.items(),
                key=lambda x: x[1],
                reverse=True
            )

            # --- 终端输出所有用户的平均置信度 ---
            print("\n=== 最近5次预测的平均置信度 ===")
            for user_idx, avg_prob in sorted_avg_probs:
                user_name = self.index_to_label[user_idx]
                print(f"用户 {user_name}: {avg_prob:.4f}")

            # 选择总概率最高的用户
            best_user_idx = max(total_probs.items(), key=lambda x: x[1])[0]
            best_user = self.index_to_label[best_user_idx]
            avg_prob = avg_probs[best_user_idx]  # 直接使用已计算的平均值

            # 更新显示
            self.last_prediction = self.current_prediction
            self.current_prediction = (best_user, avg_prob)

            return best_user, avg_prob

        except Exception as e:
            print(f"预测失败: {str(e)}")
            return None

    def get_adjusted_time(self):
        """获取调整后的时间（减去暂停时间）"""
        if self.start_time is None:
            return 0
        current_time = time.time()
        if hasattr(self, 'timer_paused') and self.timer_paused:
            return self.pause_start_time - self.start_time - getattr(self, 'paused_duration', 0)
        return current_time - self.start_time - getattr(self, 'paused_duration', 0)

    def pause_timer(self):
        """暂停计时器"""
        if not hasattr(self, 'timer_paused') or not self.timer_paused:
            self.timer_paused = True
            self.pause_start_time = time.time()

    def resume_timer(self):
        """恢复计时器"""
        if hasattr(self, 'timer_paused') and self.timer_paused:
            self.timer_paused = False
            if not hasattr(self, 'paused_duration'):
                self.paused_duration = 0
            self.paused_duration += time.time() - self.pause_start_time

    def export_key_events_to_json(self, filepath):
        """
        将按键信息导出到 JSON 文件，文件名与用户昵称相关
        :param filepath: 完整的文件路径
        """
        # 只保留前400次有效击键数据
        filtered_events = []
        valid_count = 0
        
        # 首先筛选有效事件并按时间排序
        for event in self.key_events:
            # 只记录字母、数字、空格、".,"按键
            if (event["key_content"] and 
                (event["key_content"].isalnum() or 
                event["key_content"].isspace() or 
                event["key_content"] in ',.') and
                event["absolute_press_time"] is not None and
                event["absolute_release_time"] is not None):
                
                if valid_count >= 400:
                    break
                    
                filtered_events.append({
                    "key_content": event["key_content"],
                    "absolute_press_time": event["absolute_press_time"],
                    "absolute_release_time": event["absolute_release_time"],
                    "is_correct": event["is_correct"]
                })
                valid_count += 1
        
        # 按绝对按下时间排序
        filtered_events.sort(key=lambda x: x["absolute_press_time"])
        
        # 将数据分组，每10个为一组，并计算相对时间
        grouped_keystrokes = []
        for i in range(0, min(400, len(filtered_events)), 10):
            group = filtered_events[i:i+10]
            if len(group) == 10:  # 只保存完整的组
                group_start_time = group[0]["absolute_press_time"]
                
                processed_group = []
                for j, event in enumerate(group):
                    # 计算相对于组开始时间的相对时间
                    press_time = event["absolute_press_time"] - group_start_time
                    release_time = event["absolute_release_time"] - group_start_time
                    
                    # 确保时间不为负
                    press_time = max(0, press_time)
                    release_time = max(press_time, release_time)  # 确保释放时间不小于按下时间
                    
                    processed_group.append({
                        "key_content": event["key_content"],
                        "press_time": round(press_time, 4),
                        "release_time": round(release_time, 4)
                    })
                
                grouped_keystrokes.append({
                    "subject": self.username,
                    "keystrokes": processed_group
                })

        # 将数据写入 JSON 文件
        with open(filepath, 'w', encoding='utf-8') as file:
            json.dump(grouped_keystrokes, file, indent=4)

        print(f"击键数据已保存到 {filepath} (共 {valid_count} 次有效击键，保存了 {len(grouped_keystrokes)*10} 次击键)")

    def export_existing_user_data(self, filepath="current_user_data.json"):
        """
        导出已有用户的击键数据
        :param filepath: 完整的文件路径，默认为 current_user_data.json
        """
        # 过滤有效击键数据并按时间排序
        filtered_events = []
        for event in self.key_events:
            # 只记录字母、数字、空格、符号等按键
            if (event["key_content"] and 
                (event["key_content"].isalnum() or 
                event["key_content"].isspace() or 
                event["key_content"] in ',.') and
                event["absolute_press_time"] is not None and
                event["absolute_release_time"] is not None):
                
                filtered_events.append({
                    "key_content": event["key_content"],
                    "absolute_press_time": event["absolute_press_time"],
                    "absolute_release_time": event["absolute_release_time"],
                    "is_correct": event["is_correct"]
                })
        
        # 按绝对按下时间排序
        filtered_events.sort(key=lambda x: x["absolute_press_time"])
        
        # 将数据分组，每10个为一组，并计算相对时间
        grouped_keystrokes = []
        for i in range(0, len(filtered_events), 10):
            group = filtered_events[i:i+10]
            if len(group) == 10:  # 只保存完整的组
                group_start_time = group[0]["absolute_press_time"]
                
                processed_group = []
                for event in group:
                    # 计算相对于组开始时间的相对时间
                    press_time = event["absolute_press_time"] - group_start_time
                    release_time = event["absolute_release_time"] - group_start_time
                    
                    # 确保时间不为负
                    press_time = max(0, press_time)
                    release_time = max(press_time, release_time)
                    
                    processed_group.append({
                        "key_content": event["key_content"],
                        "press_time": round(press_time, 4),
                        "release_time": round(release_time, 4)
                    })
                
                grouped_keystrokes.append({
                    "subject": self.username,
                    "keystrokes": processed_group
                })

        # 直接覆盖写入文件
        with open(filepath, 'w', encoding='utf-8') as file:
            json.dump(grouped_keystrokes, file, indent=4)

        print(f"已有用户击键数据已保存到 {filepath} (共 {len(filtered_events)} 次有效击键，保存了 {len(grouped_keystrokes)*10} 次击键)")